const express = require('express');
const axios = require('axios');
const path = require('path');

const app = express();
const PORT = 3000;
const OLLAMA_URL = 'http://192.168.1.14:11434';

app.use(express.json());
app.use(express.static('public'));

app.post('/api/chat', async (req, res) => {
    const { prompt, model = 'deepseek-r1:32b' } = req.body;
    
    if (!prompt) {
        return res.status(400).json({ error: 'Prompt is required' });
    }

    try {
        const response = await axios.post(`${OLLAMA_URL}/api/generate`, {
            model: model,
            prompt: prompt,
            // 添加性能优化参数
            options: {
                temperature: 0.8,
                top_k: 40,
                top_p: 0.9,
                repeat_penalty: 1.1,
                num_ctx: 4096,
                num_predict: 512,
                num_thread: 8
            }
        }, {
            responseType: 'stream'
        });

        res.setHeader('Content-Type', 'text/plain; charset=utf-8');
        res.setHeader('Cache-Control', 'no-cache');
        res.setHeader('Connection', 'keep-alive');

        let fullResponse = '';

        response.data.on('data', (chunk) => {
            const lines = chunk.toString().split('\n').filter(line => line.trim());
            
            for (const line of lines) {
                try {
                    const parsed = JSON.parse(line);
                    if (parsed.response) {
                        fullResponse += parsed.response;
                        res.write(parsed.response);
                    }
                    if (parsed.done) {
                        res.end();
                        return;
                    }
                } catch (e) {
                    // Skip invalid JSON lines
                }
            }
        });

        response.data.on('end', () => {
            res.end();
        });

        response.data.on('error', (error) => {
            console.error('Stream error:', error);
            res.status(500).end();
        });

    } catch (error) {
        console.error('API Error:', error.message);
        res.status(500).json({ error: 'Failed to connect to Ollama service' });
    }
});

// 添加获取可用模型的接口
app.get('/api/models', async (req, res) => {
    try {
        const response = await axios.get(`${OLLAMA_URL}/api/tags`);
        res.json(response.data);
    } catch (error) {
        console.error('Failed to fetch models:', error.message);
        res.status(500).json({ error: 'Failed to fetch models' });
    }
});

app.listen(PORT, '0.0.0.0', () => {
    console.log(`Server running at http://0.0.0.0:${PORT}`);
});
